import numpy as np
import matplotlib

matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy import signal
import time
import numpy.fft as fft


def naive_dft(x):
    """
    朴素实现的离散傅里叶变换
    时间复杂度O(N^2)
    :param x:
    :return:
    """
    N = len(x)
    X = np.zeros(N, dtype=complex)

    # 外层循环：频率索引 𝑘（相当于要计算第𝑘个频谱点）
    for k in range(N):
        # 内层循环：时域索引 𝑛,对所有时域样本求和。
        for n in range(N):
            # 把每个时域样本乘以复指数权重（旋转因子）累加到频谱点 X[k]
            X[k] += x[n] * np.exp(-2j * np.pi * k * n / N)

    return X


def fft_naive(x):
    '''
    递归实现的快速傅里叶变换（Cooley-Tukey算法）
    时间复杂度 O(N logN)
    :param x:
    :return:
    '''

    N = len(x)

    # 基本情况
    if N <= 1:
        return x

    # 检查N是否是2的幂
    if N & (N - 1) != 0:
        raise ValueError("N must be a power of 2")

    # 分离偶数和奇数索引
    # 取所有偶数索引元素（0, 2, 4...）并递归计算FFT
    even = fft_naive(x[0::2])
    # 取所有奇数索引元素（1, 3, 5...）并递归计算FFT
    odd = fft_naive(x[1::2])

    # 组合结果
    # 计算旋转因子（twiddle factors）与奇数部分的乘积
    T = [np.exp(-2j * np.pi * k / N) * odd[k] for k in range(N // 2)]
    # 组合偶数部分和奇数部分的结果
    return [even[k] + T[k] for k in range(N // 2)] + [even[k] - T[k] for k in range(N // 2)]


def plot_dft_comparison():
    """比较不同DFT实现的速度和结果"""
    # 生成测试信号
    # t是0到1秒， 256个采样点  f_s=256Hz, 采样间间隔dt=1/256
    t = np.linspace(0, 1, 256, endpoint=False)
    frequency1 = 5  # Hz
    frequency2 = 20  # Hz
    signal_data = np.sin(2 * np.pi * frequency1 * t) + 0.5 * np.sin(2 * np.pi * frequency2 * t)

    # 计算频率轴 - 修正这里的错误
    dt = t[1] - t[0]  # 采样间隔
    # 返回一个长度为N的数组，第一个N/2通常是从0到Nyquist频率，后半部分是负频率
    # 在画的时候只看正频率（frequencies[:N//2]）
    frequencies = np.fft.fftfreq(len(signal_data), dt)

    # 比较不同大小的DFT计算时间
    sizes = [16, 32, 64, 128, 256]
    naive_times = []
    fft_times = []
    numpy_times = []

    for size in sizes:
        test_signal = signal_data[:size]

        # 朴素DFT
        start = time.time()
        naive_result = naive_dft(test_signal)
        naive_times.append(time.time() - start)

        # 快速傅里叶变换 (仅对2的幂)
        if size & (size - 1) == 0:  # 检查是否是2的幂
            start = time.time()
            fft_result = fft_naive(test_signal)
            fft_times.append(time.time() - start)
        else:
            fft_times.append(np.nan)

        # NumPy FFT (参考)
        start = time.time()
        numpy_result = np.fft.fft(test_signal)
        numpy_times.append(time.time() - start)

    # 绘制性能比较
    plt.figure(figsize=(12, 8))

    # 时域信号t vs signal data
    plt.subplot(2, 2, 1)
    plt.plot(t, signal_data)
    plt.title('原始信号')
    plt.xlabel('时间 (s)')
    plt.ylabel('幅度')

    # 用np.fft.fft(signal_data) 的幅度谱，并只画正频率部分
    plt.subplot(2, 2, 2)
    fft_result = np.fft.fft(signal_data)
    # 幅值没有除以N或做单边谱转换，如果想得到真实幅值（如幅度为1的正弦对应谱线高度1）， 需要做适当缩放
    plt.plot(frequencies[:len(frequencies) // 2], np.abs(fft_result[:len(fft_result) // 2]))
    plt.title('FFT频谱')
    plt.xlabel('频率 (Hz)')
    plt.ylabel('幅度')

    plt.subplot(2, 2, 3)
    plt.plot(sizes, naive_times, 'o-', label='朴素DFT')
    valid_sizes = [sizes[i] for i in range(len(sizes)) if not np.isnan(fft_times[i])]
    valid_fft_times = [t for t in fft_times if not np.isnan(t)]
    plt.plot(valid_sizes, valid_fft_times, 's-', label='FFT')
    plt.plot(sizes, numpy_times, '^-', label='NumPy FFT')
    plt.xlabel('信号长度')
    plt.ylabel('计算时间 (秒)')
    plt.title('计算时间比较')
    plt.legend()
    plt.yscale('log')

    plt.subplot(2, 2, 4)
    # 验证我们的实现与NumPy的一致性
    test_size = 16
    test_signal = signal_data[:test_size]

    naive_result = naive_dft(test_signal)
    numpy_result = np.fft.fft(test_signal)

    plt.plot(np.abs(naive_result), 'o-', label='朴素DFT')
    plt.plot(np.abs(numpy_result), 'x-', label='NumPy FFT')
    plt.xlabel('频率索引')
    plt.ylabel('幅度')
    plt.title('结果验证 (大小=16)')
    plt.legend()

    plt.tight_layout()
    plt.show()


def demonstrate_fourier_proerties():
    """
    展示傅里叶变换的性质
    :return:
    """
    t = np.linspace(0, 1, 256, endpoint=False)
    dt = t[1] - t[0]  # 采样间隔

    # 线性性质
    signal1 = np.sin(2 * np.pi * 5 * t)
    signal2 = np.cos(2 * np.pi * 10 * t)
    combined_signal = signal1 + signal2

    fft1 = fft.fft(signal1)
    fft2 = fft.fft(signal2)
    fft_combined = fft.fft(combined_signal)
    frequencies = fft.fftfreq(len(t), dt)

    # 2.时移性质
    shifted_signal = np.roll(signal1, 50)
    fft_shifted = fft.fft(shifted_signal)

    plt.figure(figsize=(12, 10))

    # 线性性质验证
    plt.subplot(3, 2, 1)
    plt.plot(t, signal1, label='signal 1')
    plt.plot(t, signal2, label='signal 2')
    plt.plot(t, combined_signal, label='combined signal')
    plt.legend()
    plt.title("time domain signals")

    plt.subplot(3, 2, 2)
    plt.plot(frequencies[:128], np.abs(fft1)[:128], label='FFT of signal 1')
    plt.plot(frequencies[:128], np.abs(fft2)[:128], label='FFT of signal 2')
    plt.plot(frequencies[:128], np.abs(fft1 + fft2)[:128], label='FFT1+FFT2')
    plt.plot(frequencies[:128], np.abs(fft_combined)[:128], label='FFT of combined signal')
    plt.legend()
    plt.title("linear property verification")

    # 时移性质验证
    plt.subplot(3, 2, 3)
    plt.plot(t, signal1, label='original signal')
    plt.plot(t, shifted_signal, label='shifted signal')
    plt.legend()
    plt.title("time domain signals")

    plt.subplot(3, 2, 4)
    phase_original = np.angle(fft1)
    phase_shifted = np.angle(fft_shifted)
    plt.plot(frequencies[:128], phase_original[:128], label='original phase')
    plt.plot(frequencies[:128], phase_shifted[:128], label='shifted phase')
    plt.legend()
    plt.title("phase shift verification")

    # 卷积定理验证
    plt.subplot(3, 2, 5)
    kernel = np.exp(-t * 5)  # 指数衰减核
    convolved = np.convolve(signal, kernel, mode="same")
    fft_conv = fft.fft(convolved)
    fft_product = fft.fft(signal) * fft.fft(kernel)

    plt.plot(frequencies[:128], np.abs(fft_conv)[:128], label='FFT of convolved signal')
    plt.plot(frequencies[:128], np.abs(fft_product)[:128], label='FFT of product')
    plt.legend()
    plt.title("convolution theorem verification")

    plt.subplot(3, 2, 6)
    plt.plot(t, convolved, label='time domain convolved signal')
    plt.plot(t, fft.ifft(fft_product).real, '--', label='frequency domain inverse transform')
    plt.legend()
    plt.title("comparison of convolved result")

    plt.tight_layout()
    plt.show()


def understanding_frequency_axis():
    """
    理解频率轴的计算
    :return:
    """
    # 创建不同采样率的信号
    fig, axes = plt.subplot(2, 2, figsize=(12, 10))

    # 情况1： 高采样率
    fs1 = 1000  # 1000Hz采样率
    t1 = np.linspace(0, 1, fs1, endpoint=False)
    signal1 = np.sin(2 * np.pi * 50 * t1) + 0.5 * np.sin(2 * np.pi * 120 * t1)
    freq1 = fft.fftfreq(len(t1), 1 / fs1)
    fft1 = fft.fft(signal1)

    axes[0, 0].plot(t1, signal1)
    axes[0, 0].set_title(f"时域信号 (fs={fs1}Hz)")
    axes[0, 0].set_xlabel("时间(s)")
    axes[0, 0].set_ylabel("幅度")

    axes[0, 1].plot(freq1[:len(freq1) // 2], np.abs(fft1)[:len(freq1) // 2])
    axes[0, 1].set_title(f"频域信号 (fs={fs1}Hz)")
    axes[0, 1].set_xlabel("频率(Hz)")
    axes[0, 1].set_ylabel("幅度")
    axes[0, 1].axvline(x=50, color="r", linestyle="--", alpha=0.7)
    axes[0, 1].axvline(x=120, color="r", linestyle="--", alpha=0.7)

    # 情况2： 低采样率
    fs2 = 200  # 200Hz采样率
    t2 = np.linspace(0, 1, fs2, endpoint=False)
    signal2 = np.sin(2 * np.pi * 50 * t2) + 0.5 * np.sin(2 * np.pi * 120 * t2)
    freq2 = fft.fftfreq(len(t2), 1 / fs2)
    fft2 = fft.fft(signal2)

    axes[1, 0].plot(t2, signal2)
    axes[1, 0].set_title(f"时域信号 (fs={fs2}Hz)")
    axes[1, 0].set_xlabel("时间(s)")
    axes[1, 0].set_ylabel("幅度")

    axes[1, 1].plot(freq2[:len(freq2) // 2], np.abs(fft2)[:len(freq2) // 2])
    axes[1, 1].set_title(f"频域信号 (fs={fs2}Hz)")
    axes[1, 1].set_xlabel("频率(Hz)")
    axes[1, 1].set_ylabel("幅度")
    axes[1, 1].axvline(x=50, color="r", linestyle="--", alpha=0.7)
    axes[1, 1].axvline(x=120, color="r", linestyle="--", alpha=0.7)

    plt.tight_layout()
    plt.show()

    print(f"采样率 fs1={fs1},奈奎斯特频率={fs1 / 2}Hz")
    print(f"采样率 fs2={fs2},奈奎斯特频率={fs2 / 2}Hz")
    print("注意: 当信号频率超过奈奎斯特频率时会发生混叠")


if __name__ == '__main__':
    plot_dft_comparison()
